import os
from datetime import datetime
import logging

logger = logging.getLogger("AppAutomation")

def capture_screenshot(driver, test_method_name=None):
    """
    在任何地方都可以调用的截图工具函数
    
    :param driver: WebDriver实例
    :param test_method_name: 测试方法名称，用于生成截图文件名
    :return: 截图文件路径，如果失败则返回None
    """
    # 检查WebDriver实例是否为空
    if not driver:
        logger.error("WebDriver实例为空，无法截图")
        return None
    
    # 如果没有提供测试方法名称，则使用默认名称
    if not test_method_name:
        test_method_name = "error_screenshot"
    
    # 清理文件名中的非法字符
    cleaned_name = "".join(c for c in test_method_name if c.isalnum() or c in (' ', '-', '_')).rstrip()
    # 限制文件名长度
    if len(cleaned_name) > 50:
        cleaned_name = cleaned_name[:50]
    
    # 生成时间戳
    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
    # 获取项目根目录的绝对路径
    project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    screenshot_dir = os.path.join(project_root, "reports", "screenshots")
    os.makedirs(screenshot_dir, exist_ok=True)
    screenshot_path = os.path.join(screenshot_dir, f"{cleaned_name}_{timestamp}.png")
    
    # 尝试截图并保存
    try:
        driver.save_screenshot(screenshot_path)
        logger.info(f"截图已保存到: {screenshot_path}")
        return screenshot_path
    except Exception as e:
        logger.error(f"截图保存失败: {e}")
        return None